# link: https://mp.weixin.qq.com/s/NTct2_AYhz4Z8q5MYtBQcA

import backtrader as bt
import pandas as pd
import datetime
import akshare as ak

pd.set_option('display.max_columns', None)


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


codes = ['sh600276', 'sh600519', 'sh603288']

# 恒瑞医药
data1 = get_data(codes[0])
# print(data1)
# print('data1')
# 贵州茅台
data2 = get_data(codes[1])
# 海天味业
data3 = get_data(codes[2])

# %%

# 实例化策略
cerebro = bt.Cerebro()

st_date = datetime.datetime(2023, 4, 4)
ed_date = datetime.datetime(2023, 5, 11)

# 添加 600276.SH 的行情数据
datafeed1 = bt.feeds.PandasData(dataname=data1, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed1, name=codes[0])

# 添加 600519.SH 的行情数据
datafeed2 = bt.feeds.PandasData(dataname=data2, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed2, name=codes[1])

# 添加 603288.SH 的行情数据
datafeed3 = bt.feeds.PandasData(dataname=data3, fromdate=st_date, todate=ed_date)
cerebro.adddata(datafeed3, name=codes[2])


# 第1.2节：验证 line 的结构
class TestStrategy(bt.Strategy):
    def __init__(self):
        print("--------- 打印 self 策略本身的 lines ----------")
        print(self.lines.getlinealiases())
        print("--------- 打印 self.datas 第一个数据表格的 lines ----------")
        print(self.datas[0].lines.getlinealiases())
        print("--------- 打印 self.datas 第二个数据表格的 lines ----------")
        print(self.datas[1].lines.getlinealiases())
        # 计算第一个数据集的s收盘价的20日均线，返回一个 Data feed
        self.sma = bt.indicators.SimpleMovingAverage(self.datas[0].close, period=20)
        print("--------- 打印 indicators 对象的 lines ----------")
        print(self.sma.lines.getlinealiases())
        print("---------- 直接打印 indicators 对象的所有 lines -------------")
        print(self.sma.lines)
        print("---------- 直接打印 indicators 对象的第一条 lines -------------")
        print(self.sma.lines[0])

    def next(self):
        print('验证索引位置为 6 的线是不是 datetime')
        # datetime 线中的时间点存的是数字形式的时间，可以通过 bt.num2date() 方法将其转为“xxxx-xx-xx xx:xx:xx”这种形式
        print(bt.num2date(self.datas[0].lines[6][0]))


cerebro.addstrategy(TestStrategy)
result = cerebro.run()
